"""
指数数据采集器模块

负责通过akshare API获取指数数据，包括：
1. 指数基本信息
2. 指数成分股
3. 指数历史行情数据
"""
import time
import datetime
import logging
import akshare as ak
import pandas as pd
from sqlalchemy.orm import Session
from sqlalchemy.exc import IntegrityError

from modules.models.base import get_db
from modules.models.stock import SecurityInfo, DailyQuote, CollectionTask
from config.settings import DATA_COLLECTION

# 配置日志
logger = logging.getLogger("indices_collector")


class IndicesCollector:
    """指数数据采集类"""
    
    def __init__(self):
        """初始化采集器"""
        self.retry_limit = DATA_COLLECTION["retry_limit"]
        self.retry_delay = DATA_COLLECTION["retry_delay"]
        self.batch_size = DATA_COLLECTION["batch_size"]
        self.db = next(get_db())
    
    def __del__(self):
        """析构函数，确保关闭数据库连接"""
        if hasattr(self, 'db') and self.db:
            self.db.close()
    
    def get_all_indices_list(self):
        """获取所有指数列表"""
        try:
            logger.info("开始获取指数列表...")
            # 使用akshare获取指数列表
            indices_df = ak.index_stock_info()
            logger.info(f"成功获取{len(indices_df)}个指数信息")
            return indices_df
        except Exception as e:
            logger.error(f"获取指数列表失败: {str(e)}")
            raise
    
    def collect_indices_info(self):
        """采集并更新指数基本信息"""
        logger.info("开始采集指数基本信息...")
        try:
            # 获取指数列表
            indices_df = self.get_all_indices_list()
            
            # 获取现有指数代码
            existing_codes = {
                code[0] for code in 
                self.db.query(SecurityInfo.code).filter(SecurityInfo.type == 'index').all()
            }
            
            # 准备批量插入或更新的数据
            new_indices = []
            update_indices = []
            
            for _, row in indices_df.iterrows():
                code = row['code']
                name = row['name']
                exchange = row['exchange'] if 'exchange' in row else 'OTHER'
                
                # 检查是否存在
                if code in existing_codes:
                    # 更新现有指数
                    index = self.db.query(SecurityInfo).filter(
                        SecurityInfo.code == code,
                        SecurityInfo.type == 'index'
                    ).first()
                    
                    if index:
                        index.name = name
                        index.exchange = exchange
                        index.update_time = datetime.datetime.now()
                        update_indices.append(index)
                else:
                    # 创建新指数记录
                    new_index = SecurityInfo(
                        code=code,
                        name=name,
                        type='index',
                        exchange=exchange,
                        is_active=1
                    )
                    new_indices.append(new_index)
            
            # 批量保存新指数
            if new_indices:
                self.db.add_all(new_indices)
                self.db.commit()
                logger.info(f"成功添加{len(new_indices)}个新指数")
            
            # 批量更新现有指数
            if update_indices:
                self.db.commit()
                logger.info(f"成功更新{len(update_indices)}个指数信息")
            
            return len(new_indices), len(update_indices)
        
        except Exception as e:
            self.db.rollback()
            logger.error(f"采集指数基本信息失败: {str(e)}")
            raise
    
    def collect_index_components(self, index_code):
        """
        采集指数成分股
        
        Args:
            index_code: 指数代码
            
        Returns:
            成分股数量
        """
        try:
            logger.info(f"开始采集指数 {index_code} 的成分股...")
            
            # 获取指数信息
            index = self.db.query(SecurityInfo).filter(
                SecurityInfo.code == index_code,
                SecurityInfo.type == 'index'
            ).first()
            
            if not index:
                logger.error(f"指数代码 {index_code} 不存在")
                return 0
            
            # 使用akshare获取成分股
            components_df = ak.index_stock_cons(index_code)
            
            if components_df.empty:
                logger.info(f"指数 {index_code} 没有成分股数据")
                return 0
            
            # 获取现有股票代码
            existing_codes = {
                code[0] for code in 
                self.db.query(SecurityInfo.code).filter(SecurityInfo.type == 'stock').all()
            }
            
            # 更新成分股标记
            updated_count = 0
            for _, row in components_df.iterrows():
                stock_code = row['code']
                
                # 检查股票是否存在
                stock = self.db.query(SecurityInfo).filter(
                    SecurityInfo.code == stock_code,
                    SecurityInfo.type == 'stock'
                ).first()
                
                if stock:
                    # 更新成分股标记
                    stock.is_index_component = 1
                    stock.index_component_of = index_code
                    updated_count += 1
            
            if updated_count > 0:
                self.db.commit()
                logger.info(f"成功更新{updated_count}只成分股信息")
            
            return updated_count
            
        except Exception as e:
            self.db.rollback()
            logger.error(f"采集指数成分股失败: {str(e)}")
            return 0
    
    def collect_daily_quote(self, index_code, start_date=None, end_date=None, task_id=None):
        """
        采集指定指数的日线数据
        
        Args:
            index_code: 指数代码
            start_date: 开始日期，默认为None，表示从上次采集的位置继续
            end_date: 结束日期，默认为None，表示采集到当前日期
            task_id: 任务ID，用于更新任务状态
            
        Returns:
            新增数据条数
        """
        # 更新任务状态为运行中
        if task_id:
            self._update_task_status(task_id, "running")
        
        # 获取指数信息
        index = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == index_code,
            SecurityInfo.type == 'index'
        ).first()
        
        if not index:
            error_msg = f"指数代码 {index_code} 不存在"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
        
        # 如果未指定开始日期，则查找最新数据日期作为开始日期
        if not start_date:
            last_record = self.db.query(DailyQuote).filter(
                DailyQuote.security_id == index.id
            ).order_by(DailyQuote.trade_date.desc()).first()
            
            if last_record:
                # 从最后一条记录的下一天开始
                start_date = (last_record.trade_date + datetime.timedelta(days=1)).strftime('%Y%m%d')
            else:
                # 默认采集近5年数据
                start_date = (datetime.datetime.now() - datetime.timedelta(days=365*5)).strftime('%Y%m%d')
        
        # 如果未指定结束日期，则使用当前日期
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        logger.info(f"开始采集指数 {index_code} ({index.name}) 的日线数据，时间范围: {start_date} - {end_date}")
        
        try:
            # 使用akshare获取指数日线数据
            retry_count = 0
            while retry_count < self.retry_limit:
                try:
                    # 调用akshare API
                    df = ak.stock_zh_index_daily(symbol=index_code)
                    break
                except Exception as e:
                    retry_count += 1
                    logger.warning(f"采集指数 {index_code} 数据失败，重试 {retry_count}/{self.retry_limit}: {str(e)}")
                    if retry_count >= self.retry_limit:
                        raise
                    time.sleep(self.retry_delay)
            
            if df.empty:
                logger.info(f"指数 {index_code} 在指定时间范围内没有数据")
                if task_id:
                    self._update_task_status(task_id, "success", None, 0)
                return 0
            
            # 转换和清洗数据
            df = self._clean_index_data(df)
            
            # 保存到数据库
            inserted_count = self._save_daily_quotes(index.id, df)
            
            # 更新任务状态
            if task_id:
                self._update_task_status(task_id, "success", None, inserted_count)
            
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集指数 {index_code} 数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
    
    def _clean_index_data(self, df):
        """清洗和标准化指数数据"""
        if df.empty:
            return df
        
        # 标准化列名
        df.columns = [col.lower() for col in df.columns]
        
        # 确保日期列格式正确
        if 'date' in df.columns:
            df.rename(columns={'date': 'trade_date'}, inplace=True)
        
        # 标准化中文列名
        name_mapping = {
            'open': 'open',
            'close': 'close',
            'high': 'high',
            'low': 'low',
            'volume': 'volume'
        }
        
        df.rename(columns=name_mapping, inplace=True)
        
        # 处理日期格式
        if 'trade_date' in df.columns:
            df['trade_date'] = pd.to_datetime(df['trade_date'])
        
        # 处理缺失值
        numeric_cols = ['open', 'close', 'high', 'low', 'volume']
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # 删除交易量为0的无效数据
        if 'volume' in df.columns:
            df = df[df['volume'] > 0]
        
        # 保留需要的列
        keep_cols = [col for col in ['trade_date', 'open', 'high', 'low', 'close', 'volume'] if col in df.columns]
        return df[keep_cols]
    
    def _save_daily_quotes(self, security_id, df):
        """保存日线数据到数据库"""
        if df.empty:
            return 0
        
        inserted_count = 0
        
        try:
            # 获取已存在的交易日期
            existing_dates = {
                date[0] for date in 
                self.db.query(DailyQuote.trade_date).filter(
                    DailyQuote.security_id == security_id
                ).all()
            }
            
            # 准备批量插入
            batch_records = []
            
            for _, row in df.iterrows():
                trade_date = row['trade_date'].date()
                
                # 跳过已存在的记录
                if trade_date in existing_dates:
                    continue
                
                daily_quote = DailyQuote(
                    security_id=security_id,
                    trade_date=trade_date,
                    open=row.get('open'),
                    high=row.get('high'),
                    low=row.get('low'),
                    close=row.get('close'),
                    adj_close=row.get('close'),  # 指数不复权
                    volume=row.get('volume')
                )
                
                batch_records.append(daily_quote)
                inserted_count += 1
                
                # 分批提交
                if len(batch_records) >= self.batch_size:
                    self.db.add_all(batch_records)
                    self.db.commit()
                    batch_records = []
            
            # 提交剩余记录
            if batch_records:
                self.db.add_all(batch_records)
                self.db.commit()
            
            logger.info(f"成功保存 {inserted_count} 条指数日线数据")
            return inserted_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"保存指数日线数据失败(完整性错误): {str(e)}")
            return 0
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存指数日线数据失败: {str(e)}")
            raise
    
    def _update_task_status(self, task_id, status, error_message=None, result_count=None):
        """更新任务状态"""
        try:
            task = self.db.query(CollectionTask).filter(CollectionTask.task_id == task_id).first()
            if task:
                task.status = status
                if error_message:
                    task.error_message = error_message
                if result_count is not None:
                    task.result_count = result_count
                task.update_time = datetime.datetime.now()
                self.db.commit()
        except Exception as e:
            self.db.rollback()
            logger.error(f"更新任务状态失败: {str(e)}")


# 测试函数
def test_indices_collector():
    """测试指数数据采集器功能"""
    collector = IndicesCollector()
    
    # 1. 测试获取指数列表
    try:
        indices_list = collector.get_all_indices_list()
        print(f"成功获取 {len(indices_list)} 个指数基本信息")
        print(indices_list.head())
    except Exception as e:
        print(f"获取指数列表失败: {str(e)}")
    
    # 2. 测试采集指数基本信息
    try:
        new_count, update_count = collector.collect_indices_info()
        print(f"新增指数: {new_count}, 更新指数: {update_count}")
    except Exception as e:
        print(f"采集指数基本信息失败: {str(e)}")
    
    # 3. 测试采集成分股
    if indices_list is not None and not indices_list.empty:
        # 选择第一个指数进行测试
        test_code = indices_list.iloc[0]['code']
        try:
            count = collector.collect_index_components(test_code)
            print(f"指数 {test_code} 成功采集 {count} 只成分股")
        except Exception as e:
            print(f"采集指数 {test_code} 成分股失败: {str(e)}")
    
    # 4. 测试采集日线数据
    if indices_list is not None and not indices_list.empty:
        # 选择第一个指数进行测试
        test_code = indices_list.iloc[0]['code']
        try:
            # 采集最近30天数据
            end_date = datetime.datetime.now().strftime('%Y%m%d')
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')
            
            count = collector.collect_daily_quote(test_code, start_date, end_date)
            print(f"指数 {test_code} 成功采集 {count} 条日线数据")
        except Exception as e:
            print(f"采集指数 {test_code} 日线数据失败: {str(e)}")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_indices_collector()